from functools import reduce
import math
import tensorflow as tf
import tensorflow.keras as keras
import tensorflow.keras.backend as K
from tensorflow.keras import layers
from losses_utils import *


# @tf.function
def yolo_loss_on_size(labels,predns,grid,grid_anchors,
                      nc=80,na=3,label_smoothing=0.0001,positives=0.95,
                      xyxy=False,GIoU=False,DIoU=False,CIoU=False,
                      filter=True,weight=False):
    """
        predns为张量, (B,nx*ny,3,5+nc)
        labels为张量, (B,nx*ny,3,5+nc) [xc,yx,w,h,r,c...]
        #! loop 采用for循环计算, 还是张量整体计算
        #! filter 将正样本分离出来计算, 还是不分离计算
        #! weight 按正样本数量计算, 还是按权重计算
    """
    if not isinstance(labels,tf.Tensor): labels = tf.convert_to_tensor(labels,dtype=predns.dtype)
    else: labels = tf.cast(labels,dtype=predns.dtype)

    #! 初始化变量
    lbox = 0.0
    lbxy = 0.0
    lbwh = 0.0
    lcls = 0.0
    lobj = 0.0

    #! 还原pbox 从 偏移量 到 具体坐标
    pxy =  tf.nn.sigmoid(predns[:,:,:,0:2]) * 2  -  0.5     # (B,ny*nx,na,2) + (ny*nx,na,2)
    pwh = (tf.nn.sigmoid(predns[:,:,:,2:4]) * 2) ** 2 * tf.cast(grid_anchors,dtype=predns.dtype) # (B,ny*nx,na,2) * (na,2)
    predns = tf.concat([pxy,pwh,predns[...,4:]],axis=-1)

    obj_mask = labels[...,4:5] #! (B,nx*ny,na,1)
    iou = bbox_iou(labels[...,:4],predns[...,:4],xyxy=xyxy,GIoU=GIoU,DIoU=DIoU,CIoU=CIoU)[...,None] #! (B,nx*ny,na,1)
    lbox = (1-iou) * obj_mask

    lbxy = tf.stop_gradient( #! (B,nx*ny,na,1)
        K.sum(K.pow(labels[...,:2] - predns[...,:2],2),axis=-1,keepdims=True)
    )
    lbxy = lbxy * obj_mask

    lbwh = tf.stop_gradient(
        K.sum(
            K.abs(tf.nn.tanh(K.log(K.maximum(
                predns[...,2:4]/labels[...,2:4],K.epsilon()))
            )),
            axis=-1,
            keepdims=True
        )
    )
    lbwh = lbwh * obj_mask

    tobj = tf.where( #! (B,nx*ny,na,1)
        tf.equal(obj_mask,1.0),
        K.maximum(iou,0.0),
        tf.zeros_like(iou),
    )
    lobj = FocalLoss(BCEWithLogitsLoss(reduction='none'),alpha=positives,reduction='mean')(tobj,predns[...,4:5])

    lcls = FocalLoss(BCEWithLogitsLoss(reduction='none'),reduction='mean')(
        (1-label_smoothing) * labels[...,5:] + label_smoothing / nc,
        predns[...,5:],
    )

    num_pos = K.maximum(
        K.sum(K.cast(obj_mask,predns.dtype)), 1.0
    )

    lbox = K.sum(lbox) / num_pos
    lbxy = K.sum(lbxy) / num_pos
    lbwh = K.sum(lbwh) / num_pos
    lobj = K.mean(lobj)
    lcls = K.sum(lcls) / num_pos

    return lbox, lbxy, lbwh, lcls, lobj

@tf.function
def yolo_loss(labels,predns,grids,grid_anchors,
              nc=80,na=3,label_smoothing=0.0001,positives=30,
              xyxy=False,GIoU=False,DIoU=False,CIoU=False,
              filter=True,weight=False):
    """
        predns为张量, (B,nx*ny,3,5+nc)
        labels为张量, (B,nx*ny,3,5+nc) [xc,yx,w,h,r,c...]
        #! loop为False的情况下, filter失效
        #! loop为True, filter为True的情况下, weight失效
    """
    assert len(labels) == len(predns)

    lbox = 0.0
    lbxy = 0.0
    lbwh = 0.0
    lcls = 0.0
    lobj = 0.0
    obj_balance = [4.0 , 1.0, 0.4]
    bxy_balance = [0.25, 1.0, 4.0]

    for l,(label,predn) in enumerate(zip(labels,predns)):
        box_loss, bxy_loss, bwh_loss, cls_loss, obj_loss = yolo_loss_on_size(
            label,predn,grids[l],grid_anchors[l],
            nc=nc,na=na,label_smoothing=label_smoothing,positives=(1-positives)/positives,
            xyxy=xyxy,GIoU=GIoU,DIoU=DIoU,CIoU=CIoU,
            filter=filter,weight=weight)
        lbox += box_loss
        lbxy += bxy_loss
        lbxy += bxy_loss * bxy_balance[l]
        lbwh += bwh_loss
        lcls += cls_loss
        lobj += obj_loss * obj_balance[l]

    #! 使用model.compile中的loss_fn功能会把损失进行reduce化, 默认reduce函数是reduce_mean
    #! 可以使用reduce_sum, reduce_none, reduce_sum_over_batch进行替换, 但需要使用类处理
    #! 若直接使用loss_fn传入model.compile, 则会直接reduce_auto(默认是reduce_mean)
    
    return lbox, lbxy, lbwh, lcls, lobj


class YoloLoss(object):
    def __init__(self,grids,grid_anchors,nc=80,na=3,label_smoothing=0.0001,positives=30,
                 xyxy=False,GIoU=False,DIoU=False,CIoU=False,filter=True,weight=False):
        self.grids = grids
        self.grid_anchors = grid_anchors
        self.nc = nc
        self.na = na
        self.label_smoothing = label_smoothing
        self.positives = positives
        self.xyxy = xyxy
        self.GIoU = GIoU
        self.DIoU = DIoU
        self.CIoU = CIoU
        self.filter = filter
        self.weight = weight
    
    @tf.function
    def __call__(self,labels,predns):
        assert len(labels) == len(predns)

        lbox = 0.0
        lbxy = 0.0
        lbwh = 0.0
        lcls = 0.0
        lobj = 0.0
        obj_balance = [4.0 , 1.0, 0.4]
        bxy_balance = [0.25, 1.0, 4.0]

        for l,(_labels,_predns) in enumerate(zip(labels,predns)):
            _lbox, _lbxy, _lbwh, _lcls, _lobj = yolo_loss_on_size(
                _labels,_predns,self.grids[l],self.grid_anchors[l],
                nc=self.nc,na=self.na,label_smoothing=self.label_smoothing,positives=self.positives,
                xyxy=self.xyxy,GIoU=self.GIoU,DIoU=self.DIoU,CIoU=self.CIoU,
                filter=self.filter,weight=self.weight)
            
            tf.print("l: ",l)
            tf.print("lbox: ",lbox)
            tf.print("lcls: ",lcls)
            tf.print("lobj: ",lobj)     
                   
            lbox += _lbox
            lbxy += _lbxy * bxy_balance[l]
            lbwh += _lbwh
            lcls += _lcls
            lobj += _lobj * obj_balance[l]
        
        return lbox, lbxy, lbwh, lcls, lobj